"""
usage 'pinhole port host [dstport]'

Pinhole forwards the port to the host specified.
The optional dstport parameter may be used to
redirect to a different port.

eg. pinhole 80 webserver
    Forward all incoming WWW sessions to webserver.

    pinhole 23 localhost 2323
    Forward all telnet sessions to port 2323 on localhost.
"""

import sys
import os
from socket import *
from threading import Thread, Lock
import time
import struct
import traceback
import re
import random
import select

LOGGING = 1
OFFSET_LOCK = Lock()

def log( s ):
    if LOGGING:
        print '%s:%s' % ( time.ctime(), s )
        sys.stdout.flush()

class PipeThread( Thread ):
    pipes = []
    def __init__( self, mode, pad_path, offset_path, source, sink ):
        Thread.__init__( self )
        self.mode = mode
        self.pad_file = open(pad_path, 'rb')
        self.offset_file = None
        if offset:
            self.offset_path = offset_path
            #self.offset_file = open(offset, 'rb')
        self.source = source
        self.sink = sink

        log( 'Creating new pipe thread %s ( %s -> %s ) (mode %s pad_file %s offset_file %s)' % \
            ( self, source.getpeername(), sink.getpeername(), mode, pad_path, offset_path ))
        PipeThread.pipes.append( self )
        log( '%s pipes active' % len( PipeThread.pipes ))

    def run( self ):
        while 1:
            try:
                if self.mode == 'decrypt':
                    offset, length = struct.unpack( 'QQ', self.source.recv( 16 ))
                    indata = self.source.recv( length )
                    if not indata: break
                    self.pad_file.seek(offset)
                    pad = self.pad_file.read( length )
                    assert len(pad) == len(indata)
                    outdata = ''.join([chr(ord(c) ^ ord(d)) for c, d in zip(pad, indata)])
                else:
                    indata = self.source.recv( 1000 )
                    if not indata: break
                    length = len(indata)
                    OFFSET_LOCK.acquire()
                    self.offset_file = open(self.offset_path, 'rb')
                    offset, = struct.unpack('Q', self.offset_file.read(8))
                    self.offset_file.close()
                    self.offset_file = open(self.offset_path, 'wb')
                    self.offset_file.write(struct.pack('Q', offset + length))
                    self.offset_file.flush()
                    self.offset_file.close()
                    OFFSET_LOCK.release()
                    self.pad_file.seek(offset)
                    pad = self.pad_file.read( length )
                    outdata = struct.pack('QQ', offset, length ) + ''.join([chr(ord(c) ^ ord(d)) for c, d in zip(pad, indata)])
                self.sink.send( outdata )
            except:
                traceback.print_exc()
                break

        log( '%s terminating' % self )
        if self.offset_file:
            self.offset_file.close()
        self.pad_file.close()
        PipeThread.pipes.remove( self )
        log( '%s pipes active' % len( PipeThread.pipes ))

class Pinhole( Thread ):
    def __init__( self, mode, clientpad_path, serverpad_path, offset_path, protocol, myip_in, myport_in, dstip, dstport, sip_proxy=False, print_as_hex=False ):
        Thread.__init__( self )
        log( 'Redirecting: ??? -> %s:%s -> %s:%s (protocol %s mode %s out %s in %s offset %s)' % ( myip_in, myport_in, dstip, dstport, protocol, mode, serverpad_path, clientpad_path, offset_path ))
        self.mode = mode
        self.clientpad_path = clientpad_path
        self.serverpad_path = serverpad_path
        self.offset_path = offset_path
        self.protocol = protocol
        self.myip_in = myip_in
        self.myport_in = myport_in
        self.dstip = dstip
        self.dstport = dstport
        self.sip_proxy = sip_proxy
        self.print_as_hex = print_as_hex
        self.decrypt_mappings_by_name = dict()
        self.encrypt_mappings_by_name = dict()
        if protocol == 'stream':
            self.sock = socket( AF_INET, SOCK_STREAM )
            self.sock.bind(( myip_in, myport_in ))
            self.sock.listen(5)
        else:
            assert protocol == 'dgram'
            self.sockin = socket( AF_INET, SOCK_DGRAM )
            while True:
                try:
                    if myport_in is None:
                        myport_in_t = random.randrange(1000, 65000)
                    else:
                        myport_in_t = myport_in
                    self.sockin.bind(( self.myip_in, myport_in_t ))
                    myport_in = myport_in_t
                    self.myport_in = myport_in
                    break
                except:
                    if myport_in is not None:
                        raise
            self.sockout = socket( AF_INET, SOCK_DGRAM )
            self.myip_out = myip_in
            while True:
                try:
                    self.myport_out = random.randrange(1000, 65000)
                    self.sockout.bind(( self.myip_out, self.myport_out ))
                    break
                except:
                    pass
            self.destination_address = (dstip, dstport)
            self.serverpad_file = open(self.serverpad_path, 'rb')
            self.clientpad_file = open(self.clientpad_path, 'rb')

    def run( self ):
        while 1:
            if self.protocol == 'stream':
                newsock, address = self.sock.accept()
                log( 'Creating new session for %s %s ' % address )
                fwd = socket( AF_INET, SOCK_STREAM )
                fwd.connect(( self.dstip, self.dstport ))
                if self.mode == 'server':
                    PipeThread( 'decrypt', self.serverpad_path, None, newsock, fwd ).start()
                    PipeThread( 'encrypt', self.clientpad_path, self.offset_path, fwd, newsock ).start()
                else:
                    assert self.mode == 'client'
                    PipeThread( 'encrypt', self.serverpad_path, self.offset_path, newsock, fwd ).start()
                    PipeThread( 'decrypt', self.clientpad_path, None, fwd, newsock ).start()
            else:
                ready_socks, _, _ = select.select([self.sockout, self.sockin], [], [], 1000000000)
                for ready_sock in ready_socks:
                    indata, address = ready_sock.recvfrom(100000)
                    if ready_sock == self.sockin and (not hasattr(self, 'source_address') or self.source_address != address):
                        self.source_address = address
                        log( 'local address set to %r' % (address,))
                    if ready_sock == self.sockout and self.destination_address != address:
                        self.destination_address = address
                        log( 'remote address set to %r' % (address,))
                    if ((ready_sock == self.sockout and self.mode == 'client') or
                        (ready_sock == self.sockin and self.mode == 'server')):
                        if self.mode == 'client':
                            pad_file = self.clientpad_file
                            target = self.source_address
                            cursock = self.sockin
                        else:
                            assert self.mode == 'server'
                            pad_file = self.serverpad_file
                            target = self.destination_address
                            cursock = self.sockout
                        while indata:
                            offset, length = struct.unpack( 'QQ', indata[:16])
                            tindata = indata[16:16+length]
                            indata = indata[16+length:]
                            pad_file.seek(offset)
                            pad = pad_file.read( length )
                            assert len(pad) == len(tindata)
                            toutdata = ''.join([chr(ord(c) ^ ord(d)) for c, d in zip(pad, tindata)])
                            log( '-----------DECRYPT------------\n%s\n----------DECRYPT----------\n' % (toutdata.encode('hex') if self.print_as_hex else toutdata))
                            if self.sip_proxy:
                                # modify ports and ips, spawn thread to handle dgram listen add a mapping over where to forward the result
                                src_ip_match = re.search(r'c=IN IP4 ([\d\.]+)', toutdata)
                                src_port_match = re.search(r'm=audio (\d+)', toutdata)
                                if src_port_match and src_ip_match:
                                    src_ip = src_ip_match.group(1)
                                    src_port = int(src_port_match.group(1))
                                    src_address = (src_ip, src_port)
                                    mapping_name = 'audio'
                                    #mapping_name = src_address
                                    if mapping_name not in self.encrypt_mappings_by_name and mapping_name not in self.decrypt_mappings_by_name:
                                        if self.mode == 'client': # c <-o->> cg - sg - s
                                            pinhole = Pinhole( 'client', 'client.pad', 'server.pad', 'client.offset', 'dgram', self.myip_in, None, src_ip, src_port, print_as_hex=True )
                                        else: # c - cg - sg <<-o-> s
                                            assert self.mode == 'server'
                                            pinhole = Pinhole( 'client', 'server.pad', 'client.pad', 'server.offset', 'dgram', self.myip_in, None, src_ip, src_port, print_as_hex=True )
                                        pinhole.start()
                                        self.decrypt_mappings_by_name[mapping_name] = pinhole
                                    if mapping_name in self.decrypt_mappings_by_name:
                                        toutdata = toutdata.replace('IP4 ' + src_ip, 'IP4 ' + self.decrypt_mappings_by_name[mapping_name].myip_in)
                                        toutdata = toutdata.replace('m=audio ' + str(src_port), 'm=audio ' + str(self.decrypt_mappings_by_name[mapping_name].myport_in))
                                    else:
                                        toutdata = toutdata.replace('IP4 ' + src_ip, 'IP4 ' + self.encrypt_mappings_by_name[mapping_name].myip_out)
                                        toutdata = toutdata.replace('m=audio ' + str(src_port), 'm=audio ' + str(self.encrypt_mappings_by_name[mapping_name].myport_out))
                                if self.mode == 'server':
                                    toutdata = toutdata.replace(self.myip_in + ':' + str(self.myport_in), target[0] + ':' + str(target[1]))
                                    toutdata = toutdata.replace(address[0] + ':' + str(address[1]), self.myip_out + ':' + str(self.myport_out))
                                    toutdata = toutdata.replace('rport=' + str(self.myport_in) + ';', 'rport=' + str(target[1]) + ';')
                                    toutdata = toutdata.replace('rport=' + str(address[1]) + ';', 'rport=' + str(self.myport_out) + ';')
                                else:
                                    toutdata = toutdata.replace(self.myip_out + ':' + str(self.myport_out), target[0] + ':' + str(target[1]))
                                    toutdata = toutdata.replace(address[0] + ':' + str(address[1]), self.myip_in + ':' + str(self.myport_in))
                                    toutdata = toutdata.replace('rport=' + str(self.myport_out) + ';', 'rport=' + str(target[1]) + ';')
                                    toutdata = toutdata.replace('rport=' + str(address[1]) + ';', 'rport=' + str(self.myport_in) + ';')
                            log( '-----------MOD DECRYPT------------\n%s\n----------MOD DECRYPT----------\n' % (toutdata.encode('hex') if self.print_as_hex else toutdata))
                            if self.mode == 'client':
                                if self.offset_path == 'client.offset':
                                    log( '\nDETAIL: client decrypt "%r <- D(%r, %r) <- (%r, NA) <- NA" size %i\n' % (target[0] + ':' + str(target[1]), self.myip_in + ':' + str(self.myport_in), self.myip_out + ':' + str(self.myport_out), address[0] + ':' + str(address[1]), len(toutdata)))
                                else:
                                    log( '\nDETAIL: server decrypt "NA -> (NA, %r) -> D(%r, %r) -> %r" size %i\n' % (address[0] + ':' + str(address[1]), self.myip_out + ':' + str(self.myport_out), self.myip_in + ':' + str(self.myport_in), target[0] + ':' + str(target[1]), len(toutdata)))
                            if self.mode == 'server':
                                if self.offset_path == 'server.offset':
                                    log( '\nDETAIL: server decrypt "NA -> (NA, %r) -> D(%r, %r) -> %r" size %i\n' % (address[0] + ':' + str(address[1]), self.myip_in + ':' + str(self.myport_in), self.myip_out + ':' + str(self.myport_out), target[0] + ':' + str(target[1]), len(toutdata)))
                                else:
                                    log( '\nDETAIL: client decrypt "%r <- D(%r, %r) <- (%r, NA) <- NA" size %i\n' % (target[0] + ':' + str(target[1]), self.myip_out + ':' + str(self.myport_out), self.myip_in + ':' + str(self.myport_in), address[0] + ':' + str(address[1]), len(toutdata)))
                            cursock.sendto(toutdata, target)
                    elif ((ready_sock == self.sockout and self.mode == 'server') or
                          (ready_sock == self.sockin and self.mode == 'client')):
                        if self.mode == 'client':
                            pad_file = self.serverpad_file
                            target = self.destination_address
                            cursock = self.sockout
                        else:
                            assert self.mode == 'server'
                            pad_file = self.clientpad_file
                            target = self.source_address
                            cursock = self.sockin
                        log( '-----------ENCRYPT------------\n%s\n----------ENCRYPT----------\n' % (indata.encode('hex') if self.print_as_hex else indata))
                        if self.sip_proxy:
                            # modify ports and ips, spawn thread to handle dgram listen add a mapping over where to forward the result
                            src_ip_match = re.search(r'c=IN IP4 ([\d\.]+)', indata)
                            src_port_match = re.search(r'm=audio (\d+)', indata)
                            if src_port_match and src_ip_match:
                                src_ip = src_ip_match.group(1)
                                src_port = int(src_port_match.group(1))
                                src_address = (src_ip, src_port)
                                mapping_name = 'audio'
                                #mapping_name = src_address
                                if mapping_name not in self.encrypt_mappings_by_name and mapping_name not in self.decrypt_mappings_by_name:
                                    if self.mode == 'client': # c <<-o-> cg - sg - s
                                        pinhole = Pinhole( 'server', 'server.pad', 'client.pad', 'client.offset', 'dgram', self.myip_in, None, src_ip, src_port, print_as_hex=True )
                                    else: # c - cg - sg <-o->> s
                                        assert self.mode == 'server'
                                        pinhole = Pinhole( 'server', 'client.pad', 'server.pad', 'server.offset', 'dgram', self.myip_in, None, src_ip, src_port, print_as_hex=True )
                                    pinhole.start()
                                    self.encrypt_mappings_by_name[mapping_name] = pinhole
                                if mapping_name in self.decrypt_mappings_by_name:
                                    indata = indata.replace('IP4 ' + src_ip, 'IP4 ' + self.decrypt_mappings_by_name[mapping_name].myip_out)
                                    indata = indata.replace('m=audio ' + str(src_port), 'm=audio ' + str(self.decrypt_mappings_by_name[mapping_name].myport_out))
                                else:
                                    indata = indata.replace('IP4 ' + src_ip, 'IP4 ' + self.encrypt_mappings_by_name[mapping_name].myip_in)
                                    indata = indata.replace('m=audio ' + str(src_port), 'm=audio ' + str(self.encrypt_mappings_by_name[mapping_name].myport_in))
                            if self.mode == 'client':
                                indata = indata.replace(self.myip_in + ':' + str(self.myport_in), target[0] + ':' + str(target[1]))
                                indata = indata.replace(address[0] + ':' + str(address[1]), self.myip_out + ':' + str(self.myport_out))
                                indata = indata.replace('rport=' + str(self.myport_in) + ';', 'rport=' + str(target[1]) + ';')
                                indata = indata.replace('rport=' + str(address[1]) + ';', 'rport=' + str(self.myport_out) + ';')
                            else:
                                indata = indata.replace(self.myip_out + ':' + str(self.myport_out), target[0] + ':' + str(target[1]))
                                indata = indata.replace(address[0] + ':' + str(address[1]), self.myip_in + ':' + str(self.myport_in))
                                indata = indata.replace('rport=' + str(self.myport_out) + ';', 'rport=' + str(target[1]) + ';')
                                indata = indata.replace('rport=' + str(address[1]) + ';', 'rport=' + str(self.myport_in) + ';')
                        log( '-----------MOD ENCRYPT------------\n%s\n----------MOD ENCRYPT----------\n' % (indata.encode('hex') if self.print_as_hex else indata))
                        if self.mode == 'client':
                            if self.offset_path == 'client.offset':
                                log( '\nDETAIL: client encrypt "%r -> E(%r, %r) -> (%r, NA) -> NA" size %i\n' % (address[0] + ':' + str(address[1]), self.myip_in + ':' + str(self.myport_in), self.myip_out + ':' + str(self.myport_out), target[0] + ':' + str(target[1]), len(indata)))
                            else:
                                log( '\nDETAIL: server encrypt "NA <- (NA, %r) <- E(%r, %r) <- %r" size %i\n' % (target[0] + ':' + str(target[1]), self.myip_out + ':' + str(self.myport_out), self.myip_in + ':' + str(self.myport_in), address[0] + ':' + str(address[1]), len(indata)))
                        if self.mode == 'server':
                            if self.offset_path == 'server.offset':
                                log( '\nDETAIL: server encrypt "NA <- (NA, %r) <- E(%r, %r) <- %r" size %i\n' % (target[0] + ':' + str(target[1]), self.myip_in + ':' + str(self.myport_in), self.myip_out + ':' + str(self.myport_out), address[0] + ':' + str(address[1]), len(indata)))
                            else:
                                log( '\nDETAIL: client encrypt "%r -> E(%r, %r) -> (%r, NA) -> NA" size %i\n' % (address[0] + ':' + str(address[1]), self.myip_out + ':' + str(self.myport_out), self.myip_in + ':' + str(self.myport_in), target[0] + ':' + str(target[1]), len(indata)))
                        length = len(indata)
                        OFFSET_LOCK.acquire()
                        self.offset_file = open(self.offset_path, 'rb')
                        offset, = struct.unpack('Q', self.offset_file.read(8))
                        self.offset_file.close()
                        self.offset_file = open(self.offset_path, 'wb')
                        self.offset_file.write(struct.pack('Q', offset + length))
                        self.offset_file.flush()
                        self.offset_file.close()
                        OFFSET_LOCK.release()
                        pad_file.seek(offset)
                        pad = pad_file.read( length )
                        outdata = struct.pack('QQ', offset, length ) + ''.join([chr(ord(c) ^ ord(d)) for c, d in zip(pad, indata)])
                        cursock.sendto(outdata, target)

if __name__ == '__main__':

    log( 'Starting Pinhole' )

    if len( sys.argv ) >= 10:
        mode = sys.argv[1]
        clientpad = sys.argv[2]
        serverpad = sys.argv[3]
        offset = sys.argv[4]
        protocol = sys.argv[5]
        myip = sys.argv[6]
        myport = int( sys.argv[7] )
        dstip = sys.argv[8]
        dstport = int( sys.argv[9] )
        Pinhole( mode, clientpad, serverpad, offset, protocol, myip, myport, dstip, dstport ).start()
    elif len( sys.argv ) == 2:
        if sys.argv[1] == 'sip-server':
            if not os.path.exists('server.offset'):
                offset_file = open('server.offset', 'wb')
                offset_file.write(struct.pack('Q', 0))
                offset_file.close()
            sys.stdout = open( 'server.log', 'w' )
            Pinhole( 'server', 'client.pad', 'server.pad', 'server.offset', 'dgram', '192.168.0.4', 5061, '192.168.0.4', 5060, sip_proxy=True ).start()
        if sys.argv[1] == 'sip-client':
            if not os.path.exists('client.offset'):
                offset_file = open('client.offset', 'wb')
                offset_file.write(struct.pack('Q', 0))
                offset_file.close()
            sys.stdout = open( 'client.log', 'w' )
            Pinhole( 'client', 'client.pad', 'server.pad', 'client.offset', 'dgram', '192.168.0.4', 5062, '192.168.0.4', 5061, sip_proxy=True ).start()
    else:
        print "pinhole <server/client> <clientpad> <serverpad> <offset> <protocol> <myip> <myport> <dstip> <dstport> <srcip>"
        print "rdfc server.pad 10 MB"
        print "rdfc client.pad 10 MB"
        print "open('server.offset', 'wb').write(struct.pack('Q', 0))"
        print "open('client.offset', 'wb').write(struct.pack('Q', 0))"
        print "pinhole.py server client.pad server.pad server.offset dgram 127.0.0.1 5061 127.0.0.1 5060"
        print "pinhole.py server client.pad server.pad server.offset dgram 127.0.0.1 7079 127.0.0.1 7078"
        print "pinhole.py client client.pad server.pad client.offset dgram 127.0.0.1 5062 127.0.0.1 5061"
        print "pinhole.py client client.pad server.pad client.offset dgram 127.0.0.1 7080 127.0.0.1 7079 "
